from my_mol_tree import MolTree
import torch_geometric.debug as debug
import csv

def csv2txt(in_path,out_path):
    with open(in_path, 'r') as csv_file:
        csv_reader = csv.reader(csv_file)
        with open(out_path, 'w') as txt_file:
            for row in csv_reader:
                txt_file.write(row[0] + '\n')
def write_book(in_path,cli_path, edge_path):
    print("start")
    files = open(cli_path, 'w')
    i = 0
    edges_file = open(edge_path, 'w', encoding='utf-8')
    with open(in_path, "r") as f:
        for index,line in enumerate(f.readlines()):
            line = line.strip('\n')
            mol = MolTree(line)
            clique, edges = mol.get_clique_dict()
            edges_file.write(str(edges) + '\n')
            for index, cli in enumerate(clique):
                s = str(i) + ' ' + str(index) + ' /' + str(cli) + '\n'
                files.write(s)
            i += 1

    files.close()

def read_book(cli_path, edge_path):
    file = open(cli_path, 'r', encoding='utf-8')
    data = file.readlines()
    all_edge = []
    edges = open(edge_path, 'r', encoding='utf-8')
    for e in edges:
        s = e.strip('\n')
        lst = eval(s)
        # 将元组转换为二维列表
        lst_2d = [[x, y] for x, y in lst]
        all_edge.append(lst_2d)
    clique = []
    cliques = []
    mol_index = 0
    for index,cli in enumerate(data):
        cli = cli.strip('\n')
        index_cli = (cli.split(' /'))
        indexs = index_cli[0]
        mol, cli_index = indexs.split(' ')
        mol, cli_index = int(mol), int(cli_index)
        if mol != mol_index:
            cliques.append(clique)
            clique = []
            c = eval(index_cli[1])
            clique.append(c)
            mol_index=mol
        else:
            c=eval(index_cli[1])
            clique.append(c)
    cliques.append(clique)
    return cliques, all_edge



if __name__ == '__main__':
    # write_book('../data/clique_test.txt','../data/edges_test.txt')
    # all_cliques, all_edge = read_book('../data/clique_test.txt', '../data/edges_test.txt')
    # print(all_cliques)
    # print(x)
    class_name = "muv"
    root_path = "../chem_dataset/dataset/"+class_name+"/processed/"

    csv2txt(root_path+"smiles.csv",root_path+"smiles.txt")
    write_book(root_path+"smiles.txt",root_path+"clique_dict.txt",root_path+"edge_dict.txt")
